# -*- coding: utf-8 -*-
"""
@author: YuHaiyang

"""
import torch
import torchvision.models
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from torchvision.models import VGG16_Weights
from torchvision.transforms import ToTensor

from utils import pytorch_utils as utils

# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    device = utils.get_device()
    print("device:", device)

    data_train = CIFAR10(root="data", train=True, transform=ToTensor(), download=True)
    data_test = CIFAR10(root="data", train=False, transform=ToTensor(), download=True)

    dataloader_train = DataLoader(data_train, batch_size=128, shuffle=True)
    dataloader_test = DataLoader(data_test, batch_size=128)

    class_names = data_train.classes
    num_classes = len(class_names)
    print("num_classes:", num_classes)

    net = torchvision.models.vgg16(weights=VGG16_Weights.DEFAULT)
    for p in net.parameters():
        p.requires_grad = False

    net.classifier[6] = torch.nn.Linear(4096, num_classes)

    # optimizer = torch.optim.SGD(net.classifier[6], lr=0.01, momentum=0.9, weight_decay=0.005)
    optimizer = torch.optim.Adam(net.classifier[6].parameters(), lr=0.001)
    criterion = torch.nn.CrossEntropyLoss()

    net.to(device)
    num_epochs = 15
    for epoch in range(num_epochs):
        net.train()
        running_loss = 0.0
        for index, (data, label) in enumerate(dataloader_train):
            data, label = data.to(device), label.to(device)
            optimizer.zero_grad()
            out = net(data)
            loss = criterion(out, label)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        utils.save(net, "vgg16", epoch)
        print(f"Epoch {epoch + 1}/{num_epochs} Loss: {running_loss / len(dataloader_train)}")

    # net.eval()
    # correct = 0
    # total = 0
    #
    # image = Image.open("../../assets/dog.jpg")
    #
    # transform = transforms.Compose(
    #     [
    #         transforms.Resize(225),
    #         transforms.CenterCrop(224),
    #         transforms.ToTensor(),
    #         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    #     ]
    # )
    #
    # image = transform(image)
    # image = image.unsqueeze(0)
    # image = image.to(device)
    # with torch.no_grad():
    #     out_put = net(image)
    #     print("out:", out_put)
    #     pred = F.softmax(out_put[0], dim=0)
    #     print("pred:", pred)
    #
    # top5_prob, top5_catid = torch.topk(pred, 5)
    # print("top5_prob:", top5_prob)
    # print("top5_catid:", top5_catid)


    # with torch.no_grad():
    #     for inputs, labels in dataloader_test:
    #         inputs, labels = inputs.to(device), labels.to(device)
    #         outputs = net(inputs)
    #         _, predicted = torch.max(outputs.data, 1)
    #         total += labels.size(0)
    #         correct += (predicted == labels).sum().item()
    #
    # accuracy = 100 * correct / total
    # print(f"Test Accuracy: {accuracy:.2f}%")
